from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from patagent.llm import PatsnapLLM
from patagent.prompt import QUESTION_REWRITE_PRMPT


llm = PatsnapLLM()

re_write_prompt = ChatPromptTemplate.from_messages(
    [
        ("system", QUESTION_REWRITE_PRMPT),
        (
            "human",
            "Here is the initial question: \n\n {question} \n Formulate an improved question.",
        ),
    ]
)

question_rewriter = re_write_prompt | llm | StrOutputParser()
# question_rewriter.invoke({"question": question})